# coding:utf-8
import sys
import os
import re
import maya.OpenMaya as om
import maya.OpenMayaMPx as ompx
import maya.cmds as cmds

kPluginNodeTypeName = "JMaterialInfo"


class JMaterialInfo(ompx.MPxNode):
    # class variables
    TYPE_ID = om.MTypeId(0x0011A867)

    shadingGroup = om.MObject()
    surfaceShader = om.MObject()
    displacementShader = om.MObject()
    assignMeshs = om.MObject()
    outAiOpShader = om.MObject()
    outAiOpDispMap = om.MObject()

    def __init__(self):
        ompx.MPxNode.__init__(self)

    def node_name(self):
        mfn_node = om.MFnDependencyNode(self.thisMObject())
        node_name = mfn_node.name()
        return node_name

    def compute(self, plug, dataBlock):
        return
        mfn_node = om.MFnDependencyNode(self.thisMObject())
        if plug == self.outAiOpShader:
            mplug = mfn_node.findPlug(self.surfaceShader)
            if mplug.isConnected():
                input_node_array = om.MPlugArray()
                mplug.connectedTo(input_node_array, True, False)  # for get input
                input_node = om.MFnDependencyNode(input_node_array[0].node()).name()
                # 获取材质名
                out_string = "shader='{}'".format(input_node)
                # 输出给shader名
                outputHandle = dataBlock.outputValue(self.outAiOpShader)
                outputHandle.setString(out_string)
                outputHandle.setClean()
                dataBlock.setClean(plug)
        if plug == self.outAiOpDispMap:
            mplug = mfn_node.findPlug(self.displacementShader)
            if mplug.isConnected():
                input_node_array = om.MPlugArray()
                mplug.connectedTo(input_node_array, True, False)  # for get input
                input_node = om.MFnDependencyNode(input_node_array[0].node()).name()
                # 获取置换名
                out_string = "disp_map='{}'".format(input_node)
                # 输出给disp名
                outputHandle = dataBlock.outputValue(self.outAiOpDispMap)
                outputHandle.setString(out_string)
                outputHandle.setClean()
                dataBlock.setClean(plug)

        return om.kUnknownParameter

    # creator
    @staticmethod
    def nodeCreator():
        return ompx.asMPxPtr(JMaterialInfo())

    def connectionMade(self, plug, otherPlug, asSrc):
        _node = otherPlug.node()
        if plug == self.surfaceShader:
            mfn_node = om.MFnDependencyNode(self.thisMObject())
            mplug = mfn_node.findPlug(self.surfaceShader)
            if mplug.isConnected():
                input_node_array = om.MPlugArray()
                mplug.connectedTo(input_node_array, True, False)  # for get input
                input_node = om.MFnDependencyNode(input_node_array[0].node()).name()
                if input_node[0] == ':':
                    input_node = input_node[1:]
                out_string = "shader='{}'".format(input_node)
                mfn_node.findPlug(self.outAiOpShader).setString(out_string)
                self.shader_call_back_id = om.MNodeMessage.addNameChangedCallback(_node, self.shader_node_call)

        elif plug == self.displacementShader:
            mfn_node = om.MFnDependencyNode(self.thisMObject())
            mplug = mfn_node.findPlug(self.displacementShader)
            if mplug.isConnected():
                input_node_array = om.MPlugArray()
                mplug.connectedTo(input_node_array, True, False)  # for get input
                input_node = om.MFnDependencyNode(input_node_array[0].node()).name()
                if input_node[0] == ':':
                    input_node = input_node[1:]
                out_string = "disp_map='{}'".format(input_node)
                mfn_node.findPlug(self.outAiOpDispMap).setString(out_string)
                self.disp_call_back_id = om.MNodeMessage.addNameChangedCallback(_node, self.disp_node_call)
        else:
            pass
        # return om.MPxNode.connectionMade(self, plug, otherPlug, asSrc)

    def connectionBroken(self, plug, otherPlug, asSrc):
        _node = otherPlug.node()

        if plug == self.surfaceShader:
            om.MNodeMessage.removeCallback(self.shader_call_back_id)
        elif plug == self.displacementShader:
            om.MNodeMessage.removeCallback(self.disp_call_back_id)
        # return om.MPxNode.connectionBroken(self, plug, otherPlug, asSrc)

    def shader_node_call(self, plug, function, clientData):

        mfn_node = om.MFnDependencyNode(self.thisMObject())
        mplug = mfn_node.findPlug(self.surfaceShader)
        if mplug.isConnected():
            input_node_array = om.MPlugArray()
            mplug.connectedTo(input_node_array, True, False)  # for get input
            input_node = om.MFnDependencyNode(input_node_array[0].node()).name()
            if input_node[0] == ':':
                input_node = input_node[1:]
            out_string = "shader='{}'".format(input_node)
            mfn_node.findPlug(self.outAiOpShader).setString(out_string)

    def disp_node_call(self, plug, function, clientData):

        mfn_node = om.MFnDependencyNode(self.thisMObject())
        mplug = mfn_node.findPlug(self.displacementShader)
        if mplug.isConnected():
            input_node_array = om.MPlugArray()
            mplug.connectedTo(input_node_array, True, False)  # for get input
            input_node = om.MFnDependencyNode(input_node_array[0].node()).name()
            if input_node[0] == ':':
                input_node = input_node[1:]
            out_string = "disp_map='{}'".format(input_node)
            mfn_node.findPlug(self.outAiOpDispMap).setString(out_string)

    # initializer
    @staticmethod
    def nodeInitializer():
        data = om.MFnStringData()
        # input
        mAttr = om.MFnMessageAttribute()
        JMaterialInfo.shadingGroup = mAttr.create("shadingGroup", "sg")
        mAttr.setStorable(True)
        mAttr.setReadable(True)
        mAttr.setWritable(True)
        mAttr.setConnectable(True)
        mAttr.setHidden(False)

        mAttr = om.MFnMessageAttribute()
        JMaterialInfo.surfaceShader = mAttr.create("surfaceShader", "sd")
        mAttr.setStorable(True)
        mAttr.setReadable(True)
        mAttr.setWritable(True)
        mAttr.setConnectable(True)
        mAttr.setHidden(False)

        mAttr = om.MFnMessageAttribute()
        JMaterialInfo.displacementShader = mAttr.create("displacementShader", "ds")
        mAttr.setStorable(True)
        mAttr.setReadable(True)
        mAttr.setWritable(True)
        mAttr.setConnectable(True)
        mAttr.setHidden(False)

        tAttr = om.MFnTypedAttribute()
        JMaterialInfo.assignMeshs = tAttr.create("assignMeshs", "am", om.MFnData.kString, data.create(""))
        tAttr.setStorable(True)
        tAttr.setReadable(True)
        tAttr.setWritable(True)
        tAttr.setConnectable(True)
        tAttr.setHidden(False)

        # output
        tAttr = om.MFnTypedAttribute()
        JMaterialInfo.outAiOpShader = tAttr.create("outAiOpShader", "oas", om.MFnData.kString, data.create(""))
        tAttr.setHidden(0)
        tAttr.setStorable(1)
        tAttr.setWritable(0)
        tAttr = om.MFnTypedAttribute()
        JMaterialInfo.outAiOpDispMap = tAttr.create("outAiOpDispMap", "oad", om.MFnData.kString, data.create(""))
        tAttr.setHidden(0)
        tAttr.setStorable(1)
        tAttr.setWritable(0)
        #
        JMaterialInfo.addAttribute(JMaterialInfo.shadingGroup)
        JMaterialInfo.addAttribute(JMaterialInfo.surfaceShader)
        JMaterialInfo.addAttribute(JMaterialInfo.displacementShader)
        JMaterialInfo.addAttribute(JMaterialInfo.assignMeshs)
        JMaterialInfo.addAttribute(JMaterialInfo.outAiOpShader)
        JMaterialInfo.addAttribute(JMaterialInfo.outAiOpDispMap)

        JMaterialInfo.attributeAffects(JMaterialInfo.surfaceShader, JMaterialInfo.outAiOpShader)
        JMaterialInfo.attributeAffects(JMaterialInfo.displacementShader, JMaterialInfo.outAiOpDispMap)


# initialize the script plug-in
def initializePlugin(mobject):
    mplugin = ompx.MFnPlugin(mobject)
    try:
        mplugin.registerNode(kPluginNodeTypeName, JMaterialInfo.TYPE_ID, JMaterialInfo.nodeCreator,
                             JMaterialInfo.nodeInitializer)
    except:
        sys.stderr.write("Failed to register node: %s" % kPluginNodeTypeName)
        raise


# uninitialize the script plug-in
def uninitializePlugin(mobject):
    mplugin = ompx.MFnPlugin(mobject)
    try:
        mplugin.deregisterNode(JMaterialInfo.TYPE_ID)
    except:
        sys.stderr.write("Failed to deregister node: %s" % kPluginNodeTypeName)
        raise
